Which functionality not covered by existing tests does it cover?

Which functionality not covered by existing tests does it cover?

Previously our jax and pytorch distributed fused-attn only enables v2 ck backends, not v3

Which functionality not covered by existing tests does it cover?

Previously our jax and pytorch distributed fused-attn only enables v2 ck backends, not v3

Yes, but does it run different fused attn backend configs/kernels then non distributed ones? Or there is functionality concern of coexistence of them with RCCL?

Which functionality not covered by existing tests does it cover?

Previously our jax and pytorch distributed fused-attn only enables v2 ck backends, not v3

Yes, but does it run different fused attn backend configs/kernels then non distributed ones? Or there is functionality concern of coexistence of them with RCCL?

In the distributed fused-attn (CP) pytest suite, the reference run is usually a single-GPU fused-attn with full seqlen (for example sq=skv=8192) using the default attn backend. The target run decomposes the single full-size fused-attn into 4 or 8 smaller fused-attn (for example, sq=sk=4096), runs those smaller fused-attn instances using the default backend and then "glue" the results in the CP way.

In my option, why we need to enable v3 for distributed fused-attn:
1). The decomposition may create new fused-attn configs not covered by our single GPU fused-attn pytests
2). The "glue" process actually tested the softmax_lse generated in fwd pass. If the softmax lse results are not correct, the glued results will be wrong. And our single GPU fused-attn pytest do not test softmax_lse at all.

@Micky774 Could you rebase upon latest dev to incorporate the hot fix for the core sgpu tests?

@ipanfilo Could you check if all your comments have been addressed?